/// 使用epoll实现一个简单的协程框架
/// @author lcw
/// @email lcwliuchongwei@qq.com
/// 参考博客 https://zhuanlan.zhihu.com/p/401433509
use core::future::Future;
use std::{
    borrow::Borrow,
    collections::HashMap,
    mem::size_of,
    net::Ipv4Addr,
    pin::Pin,
    ptr::{addr_of, addr_of_mut, null_mut},
    str::FromStr,
    task::{Context, Poll},
};

use libc::{
    __errno_location, accept, c_void, epoll_create1, epoll_ctl, epoll_event, epoll_wait, fcntl,
    listen, setsockopt, size_t, sockaddr, sockaddr_in, socket, socklen_t, ssize_t, AF_INET, EAGAIN,
    EINTR, EPOLLET, EPOLLIN, EPOLLOUT, EPOLL_CLOEXEC, EPOLL_CTL_ADD, EPOLL_CTL_MOD, F_GETFL,
    F_SETFL, INADDR_ANY, O_NONBLOCK, SOCK_STREAM, SOL_SOCKET, SO_REUSEPORT,
};

// 封装errno宏获取错误码
macro_rules! errno {
    () => {
        unsafe { *libc::__errno_location() }
    };
}

type Task = Pin<Box<dyn Future<Output = ()>>>;

// 存储所有的异步任务
static mut TASK_ARRAY: Vec<Task> = Vec::new();

// 记录TASK_ARRAR中的空洞，即TASK_ARRAY已经完成的任务
static mut EMPTY: Vec<usize> = Vec::new();

// 记录从未被poll过的任务，第一次poll意味着被注册到epoll上去，然后变成Pending状态，等待epoll_wait唤醒
static mut NEW_TASK_STACK: Vec<usize> = Vec::new();

// 记录上次poll的fd与epoll监听的事件
// (fd, events)
static mut TASK_FD_OP: (i32, i32) = (0, 0);

pub struct TcpSocket {
    pub fd: i32,
}

impl TcpSocket {
    pub fn new() -> Self {
        unsafe {
            return Self {
                fd: socket(AF_INET, SOCK_STREAM, 0),
            };
        }
    }

    pub fn new_with_fd(fd: i32) -> Self {
        return Self { fd };
    }

    pub fn bind(&self, port: u16) -> IOResult {
        let val: i32 = 1;
        let mut sock = unsafe { std::mem::MaybeUninit::<sockaddr_in>::zeroed().assume_init() };
        //bzero
        sock.sin_family = AF_INET as u16;

        sock.sin_addr.s_addr = INADDR_ANY;
        // 将port转换成大端
        sock.sin_port = port.to_be(); // htons

        unsafe {
            setsockopt(
                self.fd,
                SOL_SOCKET,
                SO_REUSEPORT,
                addr_of!(val) as *const c_void,
                size_of::<i32>() as socklen_t,
            );
            //type SA = sockaddr;
            return if libc::bind(
                self.fd,
                addr_of!(sock) as *const sockaddr_in as *const sockaddr,
                size_of::<sockaddr_in>() as socklen_t,
            ) == -1
            {
                Err(errno!())
            } else {
                Ok(0)
            };
        }
    }

    pub fn listen(&self) -> IOResult {
        unsafe {
            let ret = if listen(self.fd, 1024) == -1 {
                Err(errno!())
            } else {
                Ok(0)
            };
            TcpSocket::set_fd_nonblocking(self.fd);

            return ret;
        }
    }

    fn set_fd_nonblocking(fd: i32) {
        unsafe {
            let flag = fcntl(fd, F_GETFL);
            fcntl(fd, F_SETFL, flag | O_NONBLOCK);
        }
    }

    pub fn accept(
        &self,
        address: Option<&mut sockaddr_in>,
        address_len: Option<&mut socklen_t>,
    ) -> AsyncAccept {
        if address.is_some() && address_len.is_some() {
            let paddr = address.unwrap() as *mut sockaddr_in as *mut sockaddr;
            let paddr_len = address_len.unwrap() as *mut socklen_t;

            return AsyncAccept::new(self.fd, paddr, paddr_len);
        }
        return AsyncAccept::new(self.fd, null_mut(), null_mut());
    }
    pub fn connect_sync<T>(&self, addr: T, port: u16) -> IOResult
    where
        T: AsRef<str>,
    {
        let mut sock = unsafe { std::mem::MaybeUninit::<sockaddr_in>::zeroed().assume_init() };

        sock.sin_family = AF_INET as u16;
        sock.sin_port = port.to_be();
        sock.sin_addr.s_addr =
            unsafe { std::mem::transmute_copy(&Ipv4Addr::from_str(addr.as_ref()).unwrap()) };

        unsafe {
            let err = libc::connect(
                self.fd,
                addr_of!(sock) as *const sockaddr_in as *const sockaddr,
                size_of::<sockaddr_in>() as u32,
            );
            return if err == 0 { Ok(0) } else { Err(errno!()) };
        }
    }

    pub fn read<B>(&self, buf: &mut B, n: usize) -> AsyncRead
    where
        B: AsMut<[u8]> + Borrow<[u8]>,
    {
        let buf = buf.as_mut();
        return AsyncRead::new(self.fd, buf, n);
    }

    pub fn write<B>(&self, buf: &B, n: size_t) -> AsyncWrite
    where
        B: AsRef<[u8]> + Borrow<[u8]>,
    {
        let buf = buf.as_ref();
        return AsyncWrite::new(self.fd, buf, n);
    }
}

pub struct Executor {
    pub epfd: i32,
    pub event_size: i32,
    pub events: Vec<epoll_event>,
    // 第一个i32表示fd，第二个i32表示监听的事件，usize表示在TASK_ARRAY的索引
    pub fd_map: HashMap<i32, (i32, usize)>,
}

impl Executor {
    pub fn new(event_size: i32) -> Self {
        return Self {
            epfd: 0,
            event_size,
            events: Vec::with_capacity(event_size as usize),
            fd_map: HashMap::new(),
        };
    }

    unsafe fn poll_new_task(&mut self, cx: &mut Context) {
        while let Some(idx) = NEW_TASK_STACK.pop() {
            let task = &mut TASK_ARRAY[idx];
            let mut ev;
            dbg!(idx, addr_of!(task));

            if task.as_mut().poll(cx).is_pending() {
                // Pending表示IO阻塞了
                self.fd_map.insert(TASK_FD_OP.0, (TASK_FD_OP.1, idx));

                ev = epoll_event {
                    events: TASK_FD_OP.1 as u32, // 存放IO事件
                    u64: TASK_FD_OP.0 as u64,    // 这个u64字段就是原本C中的data.fd
                };
                epoll_ctl(self.epfd, EPOLL_CTL_ADD, TASK_FD_OP.0, addr_of_mut!(ev));
            } else {
                // Ready表示IO操作结束可以关闭连接
                EMPTY.push(idx);
            }

            dbg!(idx, addr_of!(task), "end of poll");
        }
    }

    pub fn run(&mut self) -> ! {
        // let waker = Waker::from_raw(RawWaker::new(null(), get_raw_waker_vtable()));
        // let mut cx = Context::from_waker(&waker);
        // 创建一个没有用处的noop_waker
        let waker = futures::task::noop_waker();
        let mut cx = Context::from_waker(&waker);

        // 实例化epoll
        self.epfd = unsafe { epoll_create1(EPOLL_CLOEXEC) };

        loop {
            // n表示当前epoll_wait后有多少个就绪事件可以处理
            let n = unsafe {
                // 尝试poll新的任务
                self.poll_new_task(&mut cx);
                // epoll_wait永久阻塞，只要有事件的时候我们才需要处理
                let n = epoll_wait(self.epfd, self.events.as_mut_ptr(), self.event_size, -1);
                n
            };
            if n == -1 && errno!() == EINTR {
                // 如果出现EINTR错误继续即可，EINTR表示读写中断错误
                continue;
            }
            for i in 0..(n as usize) {
                let ev_ptr = self.events.as_ptr();
                /*
                    epoll_wait系统调用只会更改self.events.as_mut_ptr()实际值。
                    但是不会更改len导致使用rust api访问会出错，可以考虑更换nix里面封装了安全的api
                */
                let fd = unsafe { (*ev_ptr.offset(i as isize)).u64 as i32 };
                let (event, idx) = self.fd_map[&fd].clone();
                unsafe {
                    dbg!(idx, addr_of!(TASK_ARRAY[idx]));
                    // 如果idx在这次poll中完成了IO操作，就将他加入到EMPTY中
                    if TASK_ARRAY[idx].as_mut().poll(&mut cx).is_ready() {
                        EMPTY.push(idx);
                        continue;
                    }
                    dbg!(idx, addr_of!(TASK_ARRAY[idx]), "end of poll");
                    // 能走到这里表示这次poll IO操作还没有完成
                    if fd != TASK_FD_OP.0 {
                        // 这次poll中又发生了新的IO操作，需要基础epoll_ctl注册
                        let mut ev = epoll_event {
                            u64: TASK_FD_OP.0 as u64,
                            events: TASK_FD_OP.1 as u32,
                        };
                        epoll_ctl(self.epfd, EPOLL_CTL_ADD, TASK_FD_OP.0, addr_of_mut!(ev));
                        self.fd_map.insert(TASK_FD_OP.0, (TASK_FD_OP.1, idx));
                    }
                    // event与当前poll的event不一样表示修改了event，我们只需要通过epoll_ctl修改即可
                    if event != TASK_FD_OP.1 {
                        let mut ev = epoll_event {
                            u64: TASK_FD_OP.0 as u64,
                            events: TASK_FD_OP.1 as u32,
                        };
                        epoll_ctl(self.epfd, EPOLL_CTL_MOD, TASK_FD_OP.0, addr_of_mut!(ev));
                        // 将fd的event更新为新的event
                        self.fd_map.get_mut(&fd).unwrap().0 = TASK_FD_OP.1;
                    }
                }
            }
        }
    }
}

// 分配新的异步任务
pub fn spawn<F>(future: F)
where
    F: Future<Output = ()> + 'static,
{
    unsafe {
        // EMPTY数组记录TASK_ARRAY中的空位（即有task已经完成了），如果有空洞就将任务放到空位上
        // 然后将新的任务即从未被poll过的任务的索引加入NEW_TASK_STACK中
        if let Some(idx) = EMPTY.pop() {
            TASK_ARRAY[idx] = Box::pin(future);
            NEW_TASK_STACK.push(idx);
        } else {
            TASK_ARRAY.push(Box::pin(future));
            NEW_TASK_STACK.push(TASK_ARRAY.len() - 1);
        }
    }
}

pub struct AsyncAccept {
    pub fd: i32,
    pub addr: *mut sockaddr,
    pub addr_len: *mut socklen_t,
}

impl AsyncAccept {
    pub fn new(fd: i32, addr: *mut sockaddr, addr_len: *mut socklen_t) -> Self {
        return Self { fd, addr, addr_len };
    }
}

type IOResult = Result<ssize_t, i32>;

impl Future for AsyncAccept {
    type Output = IOResult;

    fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Self::Output> {
        let this = self.get_mut();
        let fd;

        // 尝试接受连接，注意这里最好是非阻塞的，创建socket时就要设置非阻塞
        fd = unsafe { accept(this.fd, this.addr, this.addr_len) };
        /*
        在非阻塞模式下，如果有错但是是EAGAIN，说明远端已经没有数据需要传输了。
        我们不需要阻塞，而是继续处理其他的fd。因为后面还要处理该fd所以不要关闭连连接
        */
        if fd == -1 && errno!() == EAGAIN {
            unsafe { TASK_FD_OP = (this.fd, EPOLLIN) };
            return Poll::Pending;
        }

        return Poll::Ready(if fd > 0 {
            Ok(fd as isize)
        } else {
            Err(errno!())
        });
    }
}

pub struct AsyncRead {
    fd: i32,
    buf: *mut c_void,
    n: size_t,
    first: bool,
}

impl AsyncRead {
    pub fn new(fd: i32, buf: &mut [u8], n: usize) -> Self {
        return Self {
            fd,
            buf: buf.as_mut_ptr() as *mut c_void,
            n,
            first: true,
        };
    }
}

impl Future for AsyncRead {
    type Output = IOResult;

    fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Self::Output> {
        let this = self.get_mut();
        let n;
        unsafe { TASK_FD_OP = (this.fd, EPOLLIN) };

        // 如果是一次poll直接返回Pending即可，接下里听从调度器安排就行
        if this.first {
            this.first = false;
            return Poll::Pending;
        }

        n = unsafe { libc::read(this.fd, this.buf, this.n) };

        return if n >= 0 {
            Poll::Ready(Ok(n))
        } else {
            Poll::Ready(Err(errno!()))
        };
    }
}

pub struct AsyncWrite {
    fd: i32,
    buf: *const c_void,
    n: size_t,
}

impl AsyncWrite {
    pub fn new(fd: i32, buf: &[u8], n: size_t) -> Self {
        return Self {
            fd,
            n,
            buf: buf.as_ptr() as *const c_void,
        };
    }
}

impl Future for AsyncWrite {
    type Output = IOResult;

    fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Self::Output> {
        let this = self.get_mut();

        unsafe {
            let n = libc::write(this.fd, this.buf, this.n);
            let code = errno!();

            // 如果出现EAGAIN错误，再试一次即可，读写操作的内核缓冲区没满是不会出现问题的
            if n == -1 && code == EAGAIN {
                TASK_FD_OP = (this.fd, (EPOLLOUT | EPOLLET));
                // TASK_FD_OP = (this.fd, (EPOLLOUT));

                return Poll::Pending;
            }

            return if n >= 0 {
                Poll::Ready(Ok(n))
            } else {
                Poll::Ready(Err(code))
            };
        }
    }
}

fn main() {
    /*
    todo 有bug调度器会poll已经完成的Future导致panic
     */
    let listener = TcpSocket::new();
    let mut ex = Executor::new(20);

    spawn(async move {
        listener.bind(8000).unwrap();
        listener.listen().unwrap();

        loop {
            println!("accepting...");
            match listener.accept(None, None).await {
                Ok(fd) => {
                    println!("accept ok");

                    spawn(async move {
                        let sock = TcpSocket::new_with_fd(fd as i32);
                        let mut buf = [0u8; 4096];

                        loop {
                            match sock.read(&mut buf, 4096).await {
                                Ok(n) => {
                                    if n == 0 {
                                        break;
                                    }

                                    println!(
                                        "read {} bytes from client, content: {}",
                                        n,
                                        String::from_utf8_lossy(&buf[0..(n as usize)])
                                    );
                                    let msg = "hello world";
                                    for (idx, c) in msg.bytes().enumerate() {
                                        buf[idx] = c;
                                    }
                                    match sock.write(&buf, msg.len()).await {
                                        Ok(n) => {
                                            println!(
                                                "write {} bytes to client, content: {}",
                                                n,
                                                String::from_utf8_lossy(&buf[0..(n as usize)])
                                            );
                                        }

                                        Err(code) => {
                                            println!("write err: {}", code);
                                            break;
                                        }
                                    }
                                }

                                Err(code) => {
                                    println!("read err: {}", code);
                                    break;
                                }
                            }
                        }
                    });
                }

                Err(code) => {
                    println!("accept err: {}", code);
                }
            }
        }
    });
    ex.run();
}
